   
# ---------------------------------------------------------------
## Scripts to load CITUP and PyClone to plot trees (Figure 2)
## Author: Veronique LeBlanc
## https://github.com/vleblanc/GBM-PDO-paper/blob/main/WES/fig2_trees.R
## https://github.com/hanhanial/PyClone_VI/blob/c2b7af0eb01d2f1e5fcadf907948bec3b5cb98d1/scripts/S05_plotting_citup-clone-trees.r
# ---------------------------------------------------------------

library(igraph)
library(rhdf5)
library(ggplot2)
library(argparser)
library(reshape2)
library(data.table)
library(dplyr)
library(ggraph)
library(tidygraph)
library(gridExtra)

##############################################################################

argp <- arg_parser("Plot the CitupPlot")
argp <- add_argument(argp, "--work_dir" , help="")
argp <- add_argument(argp, "--citup_file" , help="")
argp <- add_argument(argp, "--pyclone_file" , help="")
argp <- add_argument(argp, "--smg_gene_file" , help="")
argp <- add_argument(argp, "--citup_input_path" , help="")
argp <- add_argument(argp, "--sample" , help="")
argp <- add_argument(argp, "--output_path" , help="")
argp <- add_argument(argp, "--class_order" , help="")
argp <- add_argument(argp, "--class_a_order" , help="")
argp <- add_argument(argp, "--sample_file" , help="")

argv <- parse_args(argp)

work_dir <- argv$work_dir
citup_file <- argv$citup_file
pyclone_file <- argv$pyclone_file
smg_gene_file <- argv$smg_gene_file
citup_input_path <- argv$citup_input_path
gene_list <- argv$gene_list
sample <- argv$sample
output_path <- argv$output_path
class_order <- argv$class_order
class_a_order <- argv$class_a_order
sample_file <- argv$sample_file

if(1!=1){

	work_dir <- "~/20220915_gastric_multiple/dna_combinePublic"
	citup_file <- paste0( work_dir , "/Pyclone/Citeup_subClone_result/S63_Citeup.h5")
	pyclone_file <- paste0( work_dir , "/Pyclone/result/CiteUp/S63/tables/loci.tsv")
	citup_input_path <- paste0( work_dir , "/Pyclone/Citeup_subClone")
	smg_gene_file <- paste0( work_dir , "/public_ref/importTantGene.list")
	sample <- "S63"
	output_path <- paste0( work_dir , "/Pyclone/Citeup_subClone_plot")
  class_order <- paste(work_dir,"/config/Class_order.list",sep="")
  class_a_order <- paste(work_dir,"/config/Class_order_sub.list",sep="")
  sample_file <- paste(work_dir,"/config/tumor_normal.class.MSS_MSI.list",sep="")
}

dir.create(output_path , recursive = T)

##############################################################################

sample_id_file <- paste0( citup_input_path , "/" , sample , "_sampleID.tsv" )
cluster_file <- paste0( citup_input_path , "/" , sample , "_cluster.tsv" )
freq_file <- paste0( citup_input_path , "/" , sample , "_freq.tsv" )

# Load SampleID
sample_id <- data.frame(fread(sample_id_file , header = F))$V1
sample_id <- gsub("[.]" , "-" ,sample_id)

# 导入cluster
# 只纳入最后分析的cluster
use_cluster <- data.frame(fread(cluster_file))
use_cluster <- unique(use_cluster$V1)

## 需标记在树上的基因
dat_smg <- data.frame(fread(smg_gene_file))

geneShow_list <- unique(c(dat_smg$Gene_Symbol))

## 导入freq，后面用来做cluster的ccf在样本间的分布
dat_freq <- data.frame(fread(freq_file))

##############################################################################

class_order <- data.frame(fread(class_order))
class_a_order <- data.frame(fread(class_a_order))
info <- data.frame(fread(sample_file))

plot_id <- unique(info[info$Normal==sample,"ID"])

##############################################################################

Variant_Type <- c("Missense_Mutation","Nonsense_Mutation","Frame_Shift_Ins","Frame_Shift_Del","In_Frame_Ins","In_Frame_Del","Splice_Site","Nonstop_Mutation")

# Colormap from stackexchange, good for class palette
c25 = c("dodgerblue2","#E31A1C", # red
        "green4",
        "#6A3D9A", # purple
        "#FF7F00", # orange
        "gold1",
        "skyblue2","#FB9A99", # lt pink
        "palegreen2",
        "#CAB2D6", # lt purple
        "#FDBF6F", # lt orange
        "gray70", "khaki2",
        "maroon","orchid1","deeppink1","blue1","steelblue4",
        "darkturquoise","green1","yellow4","yellow3",
        "darkorange4","brown")

##############################################################################
# Load CITUP results
citup_iter <- h5dump(citup_file)

##############################################################################

# Get optimal solution
citup_iter_sol <- as.character(citup_iter$results$optimal$index[1])

# Get clone info from optimal solution
citup_iter_clones <- setNames(reshape2::melt(citup_iter$trees[[citup_iter_sol]]$clone_freq$block0_values),
                                    c("clone_id", "sample_id", "clonal_prev"))
citup_iter_clones$sample_id <- sample_id[citup_iter_clones$sample_id]
citup_iter_clones$clone_id <- citup_iter_clones$clone_id - 1


##############################################################################
# Load PyClone loci-level results
pyclone_loci <- read.delim(pyclone_file, stringsAsFactors = FALSE)
pyclone_loci$cluster_id <- pyclone_loci$cluster_id

pyclone_loci$sample_id <- factor(pyclone_loci$sample_id,
                                       levels = sample_id , order = T)

pyclone_loci <- subset(pyclone_loci , cluster_id %in% use_cluster )
pyclone_loci$Variant_Classification <- sapply( strsplit(pyclone_loci$mutation_id,":") , "[" , 2 )
pyclone_loci$gene <- sapply( strsplit(pyclone_loci$mutation_id,":") , "[" , 1 )


##############################################################################
# Add clone info
# citup的克隆和pyclone的对应起来

## clone的id
clone_id <- citup_iter$trees[[citup_iter_sol]]$cluster_assignment$values
names(clone_id) <- citup_iter$trees[[citup_iter_sol]]$cluster_assignment$index
clone_id2 <- clone_id

## cluster的id
pyclone_loci$cluster_id <- as.factor(pyclone_loci$cluster_id)

pyclone_loci$clone_id <- clone_id[as.character(pyclone_loci$cluster_id)]
pyclone_loci$clone_id <- as.factor(pyclone_loci$clone_id)
## clone_id和cluster_id对应
clone_cluster <- unique(data.frame(cluster_id = pyclone_loci$cluster_id , clone_id = pyclone_loci$clone_id))
clone_cluster$cluster_id <- as.numeric(as.character(clone_cluster$cluster_id))

## 一个clone会对应多个cluster
clone_cluster2 <- clone_cluster %>%
group_by(clone_id) %>%
summarize(cluster_id = paste0(cluster_id , collapse = ","))
clone_cluster2 <- data.frame(clone_cluster2)



##############################################################################
## 得到driver突变的信息
## 看每个cluster有哪些driver的突变
#driver_annotation <- subset(pyclone_loci , gene %in% geneShow_list & Variant_Classification %in% Variant_Type)
driver_annotation <- subset(pyclone_loci , Variant_Classification %in% Variant_Type)
driver_annotation <- driver_annotation %>% 
group_by(clone_id) %>%
summarize(mutation = paste0(unique(gene) , collapse = ",\n") , driver_num = length(unique(gene)) )
driver_annotation_clone_id <- as.character(driver_annotation$clone_id)

driver_annotation$clone_id <- as.character(driver_annotation$clone_id)
for(clone_id in unique(clone_cluster$clone_id)){
	if(!(clone_id %in% driver_annotation_clone_id)){
		tmp <- data.frame(clone_id = clone_id , mutation = "" , driver_num = 0 )
		driver_annotation <- rbind(driver_annotation ,tmp)
	}
}

driver_annotation$clone_id <- as.numeric(driver_annotation$clone_id)
driver_annotation$clone_id <- as.factor(driver_annotation$clone_id)
driver_annotation <- data.frame(driver_annotation)

##############################################################################
## 每个克隆突变的数目
clone_size <- pyclone_loci %>%
group_by(clone_id) %>%
summarize(num_muts = length(unique(mutation_id)))
clone_size <- data.frame(clone_size)


##############################################################################
# For igraph, edge starts from 1. Citup output starts from 0.
# I manually add 1 so it starts from 1. Then add a pseudo root node 1 to draw a
# line to indicate root.
adjacency = data.frame(t(citup_iter$trees[[citup_iter_sol]]$adjacency_list$block0_values))
colnames(adjacency) <- c("V1" , "V2")


if(length(grep( 0 , clone_cluster$clone_id)) > 0){
  ## 存在clone被判断为根克隆
  adjacency = rbind(c(-1,0), adjacency)

  adjacency_igraph = graph_from_edgelist(as.matrix(adjacency) + 2)
  ## 存在一个cluster对应多个clone
  dupclone <- names(which(duplicated(clone_id2)))
  if(length(dupclone) > 0){
    dropclone <- names(which(clone_id2==clone_id2[dupclone])[1])

    adjacency[which(adjacency$V2==dropclone),"V2"] <- as.numeric(dupclone)
    adjacency[which(adjacency$V1==dropclone),"V1"] <- as.numeric(dupclone)
    adjacency <- subset(adjacency , V1 != V2)

    adjacency_igraph = graph_from_edgelist(as.matrix(adjacency) + 2)
  }

  
}else{
  ## 不存在clone被判断为根克隆
  adjacency_igraph = graph_from_edgelist(as.matrix(adjacency) + 1)
}


# annotate driver mutations on edge
E(adjacency_igraph)$mutations = driver_annotation[match(adjacency$V2, driver_annotation$clone_id),"mutation"]

# Set vertices' names
V(adjacency_igraph)$Clones = as.character(min(adjacency):max(adjacency))

# Set root node name to none
V(adjacency_igraph)$Clones[1] = ""

# get number of driver muts for each node
V(adjacency_igraph)$driver_num = driver_annotation[match(V(adjacency_igraph)$Clones, 
                                                         driver_annotation$clone_id), 'driver_num']
V(adjacency_igraph)$driver_num[is.na(V(adjacency_igraph)$driver_num)] = 0

# number of mutations in each clone
V(adjacency_igraph)$num_muts = clone_size$num_muts[match(V(adjacency_igraph)$Clones, 
                                                         clone_size$clone_id)]
V(adjacency_igraph)$num_muts[is.na(V(adjacency_igraph)$num_muts)] = 0


# 注释clusterID
## 一个clone对应一个cluster
if(length(grep( "," , clone_cluster2$cluster_id))==0){
  V(adjacency_igraph)$cluster_id = clone_cluster$cluster_id[match(V(adjacency_igraph)$Clones, 
                                                         clone_cluster$clone_id)]
}else{
  V(adjacency_igraph)$cluster_id = clone_cluster2$cluster_id[match(V(adjacency_igraph)$Clones, 
                                                         clone_cluster2$clone_id)]
}

V(adjacency_igraph)$cluster_id[is.na(V(adjacency_igraph)$cluster_id)] = ""

##############################################################################
adjacency_tbl = as_tbl_graph(adjacency_igraph)

graph_plot = ggraph(adjacency_tbl, 'tree') +
  # Use white for the pseudo root node to hide it.
  scale_colour_manual(values = c("grey", c25[2:length(V(adjacency_igraph))])) +
  geom_edge_link(
  	aes(label=mutations), hjust=1.2, show.legend = FALSE , edge_width = 1.5,
    label_alpha=1,label_size = 3, angle_calc = "none" ) +
  # size of the nodes (clones) is proportional to number of muts in that clone
  geom_node_point(aes(size = num_muts , color=Clones), 
                  show.legend = c(size=FALSE,
                                  color=FALSE,
                                  alpha=FALSE)) +
  # 添加node的文字
  geom_node_text( aes(label = cluster_id ) , colour = "white", size = 4) +
  # Scale point to make bigger as some nodes don't have a lot of muts
  scale_size_continuous(range = c(9,15)) +
  ggtitle(plot_id) +
  theme_graph(foreground = 'steelblue', fg_text_colour = 'white',
              base_family = "sans") +
  theme(plot.title = element_text(size = 18,hjust = 0.5),
        panel.border = element_blank())


# Get x and y range to give space to labels
xlims = ggplot_build(graph_plot)$layout$panel_scales_x[[1]]$range$range
xlims[1] = xlims[1] - 0.25
xlims[2] = xlims[2] + 0.25

graph_plot = graph_plot + xlim(xlims)

images_name <- paste0(output_path , "/" , sample , "_plot_igraph.pdf")
ggsave(images_name , graph_plot)

##############################################################################
# Get proportions of malignant cells assigned to each clone
dat_ccf <- dat_freq
colnames(dat_ccf) <- sample_id
dat_ccf$cluster_id <- use_cluster

dat_ccf <- melt(data.table(dat_ccf),id.vars=c("cluster_id"))

##############################################################################
## 颜色与亚克隆树保持一致
col_use <- c25[2:length(V(adjacency_igraph))]
names(col_use) <- min(as.numeric(as.character(clone_cluster$clone_id))):max(as.numeric(as.character(clone_cluster$clone_id)))

## 颜色与亚克隆树保持一致
col_cluster <- c()
for(clone in names(col_use)){
  cluster_num <- length(grep( clone , clone_cluster$clone_id))
  cluster_ids <- clone_cluster[grep( clone , clone_cluster$clone_id),"cluster_id"]

  ## clone与cluster为一一对应
  if(cluster_num==1){
    tmp_col <- col_use[clone]
    names(tmp_col) <- cluster_ids
    col_cluster <- c( col_cluster , tmp_col)

  }else{
    ## clone与cluster为一多对应
    ## 第一个cluster与clone一样的颜色，其它cluster选择其它颜色
    tmp_col <- col_use[clone]
    names(tmp_col) <- cluster_ids[1]
    col_cluster <- c( col_cluster , tmp_col)

    for(i in 2:cluster_num){
      tmp_col <- c25[length(c25)-i]
      names(tmp_col) <- cluster_ids[i]
      col_cluster <- c( col_cluster , tmp_col)
    }
   
  }
}

#####################################################################################################
## cluster与clone对应 
dat_ccf <- merge(dat_ccf , clone_cluster , by="cluster_id")
dat_ccf <- merge(dat_ccf , clone_size , by = "clone_id")

## 提取class
dat_ccf$Class <- gsub("-." , "" , dat_ccf$variable)
dat_ccf <- data.frame(dat_ccf)

##
dat_ccf$value <- as.numeric(dat_ccf$value)
dat_ccf$cluster_id <- factor(dat_ccf$cluster_id)
dat_ccf$Class <- factor(dat_ccf$Class , levels = class_order$Class ,order = T )

##############################################################################

images_name <- paste0(output_path , "/" , sample , "_plot_constituen.pdf")

p <- ggplot(data = dat_ccf ) + 
		geom_line(aes(x = Class, y = value  , group = factor(cluster_id) ,  color = cluster_id  ) , linetype="dashed" ) +
		geom_point( aes( x = Class, y = value  , group = factor(cluster_id) , size = num_muts , color = cluster_id  )) +
		ylab("Cellular prevelance") + 
    xlab("") + 
		ylim(0 , 1) +
		scale_color_manual(values=col_cluster) +
		theme_bw() +
		#ggtitle(plot_id) +
  		theme(panel.background = element_blank(),#设置背影为白色#清除网格线
        legend.position ='right',
        legend.box = "vertical" ,
        panel.grid.major=element_line(colour=NA),
        legend.text = element_text(size = 8,color="black",face='bold'),
        axis.text.x = element_text(size = 10,color="black",face='bold'),
        axis.text.y = element_text(size = 10,color="black",face='bold'),
        axis.title.x = element_text(size = 12,color="black",face='bold'),
        axis.title.y = element_text(size = 12,color="black",face='bold'),
        axis.line = element_line(size = 0.5)) 

ggsave(images_name , p  )

##############################################################################

images_name <- paste0(output_path , "/" , sample , "_plot_combine.pdf")

plot <- grid.arrange(graph_plot , p ,ncol = 2)

ggsave(images_name , plot ,  width = 13 )
